import os
import json
import time
import torch
import numpy as np
import logging
from datetime import datetime
from mqttRev import SimpleMQTT
from test8 import WearModel, preprocess_channel

# ============ 配置 ============
RECV_TOPIC = "/iios/coco/mqtt-send/session/04f5c05f-db68-4ec4-9368-ed0ef9c60dda/msg"
SEND_TOPIC = "/iios/coco/mqtt-send/session/86e92aaa-1e6c-4eb3-823e-8a018e88a288/msg"

MODEL_PATH = "./logs/best_model1_20251031_160443.pth"  # 修改为你的模型路径
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
MAX_LEN = 32000
USE_CHANNELS = ["ch11", "ch12", "ch13", "ch14", "ch21", "ch22", "ch23"]

# ============ 日志配置 ============
os.makedirs("logs", exist_ok=True)
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
log_file = os.path.join("logs", f"mqtt_predictor_{timestamp}.log")

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    handlers=[
        logging.FileHandler(log_file, encoding="utf-8"),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger("mqtt_predictor")

# ============ 初始化 MQTT ============
mqtt_client = SimpleMQTT(
    host="192.168.6.7",
    port=2771,
    username="edge8563",
    password="XnZtSd1WlvF7Qt2E"
)

# ============ 加载模型 ============
try:
    model = WearModel().to(DEVICE)
    model.load_state_dict(torch.load(MODEL_PATH, map_location=DEVICE))
    model.eval()
    logger.info(f"✅ 模型已加载: {MODEL_PATH} (device={DEVICE})")
except Exception as e:
    logger.critical(f"❌ 模型加载失败: {e}", exc_info=True)
    raise SystemExit(1)

# ============ 数据缓存 ============
channel_buffer = {ch: [] for ch in USE_CHANNELS}


def try_predict():
    """当所有通道数据够长时执行一次预测"""
    global channel_buffer

    if all(len(channel_buffer[ch]) >= MAX_LEN for ch in USE_CHANNELS):
        try:
            data = []
            for ch in USE_CHANNELS:
                arr = np.array(channel_buffer[ch][:MAX_LEN])
                arr = preprocess_channel(arr)
                data.append(arr)
                # 清除已用部分
                channel_buffer[ch] = channel_buffer[ch][MAX_LEN:]

            data = np.array(data)
            input_tensor = torch.FloatTensor(data).unsqueeze(0).to(DEVICE)

            start_time = time.time()
            with torch.no_grad():
                pred = model(input_tensor).item()
            elapsed = time.time() - start_time

            logger.info(f"✅ 预测完成: 值={pred:.4f}, 用时={elapsed:.3f}s")

            payload = {
                "device": "ToolWearPredictor",
                "timestamp": time.time(),
                "prediction": round(pred, 4)
            }
            mqtt_client.publish(payload, topic=SEND_TOPIC)
            logger.info(f"📤 已通过 MQTT 发送预测结果: {payload}")
            time.sleep(10)
        except Exception as e:
            logger.error(f"❌ 预测阶段出错: {e}", exc_info=True)


def on_message(topic, payload):
    """MQTT 消息回调"""
    try:
        # 打印部分原始消息（防止日志太长）
        logger.info(f"📩 收到 MQTT 消息: {payload[:150]}...")

        msg = json.loads(payload)
        # 如果是数组格式，取第一个元素
        if isinstance(msg, list) and len(msg) > 0:
            msg = msg[0]

        # 更新通道数据
        for ch in USE_CHANNELS:
            if ch in msg and isinstance(msg[ch], list):
                channel_buffer[ch].extend(msg[ch])

        lengths = {ch: len(channel_buffer[ch]) for ch in USE_CHANNELS}
        logger.info(f"📥 当前各通道长度: {lengths}")

        try_predict()

    except json.JSONDecodeError as e:
        logger.warning(f"⚠️ JSON 解析失败: {e}")
    except Exception as e:
        logger.error(f"❌ on_message 异常: {e}", exc_info=True)


def main():
    mqtt_client.wait_connected(timeout=5)
    mqtt_client.subscribe(RECV_TOPIC, on_message)
    logger.info("🚀 MQTT 预测服务已启动，等待数据中...")

    try:
        while True:
            time.sleep(0.5)
    except KeyboardInterrupt:
        mqtt_client.stop()
        logger.info("🛑 用户中断，程序已退出。")


if __name__ == "__main__":
    main()
